# app.py
from flask import Flask, request, render_template, send_file,jsonify,redirect, url_for, flash
from werkzeug.utils import secure_filename
from flask_login import current_user
from utils.data_decryptor import decrypt_file
import pandas as pd
import os
import re
import hashlib
import random
from Crypto.Cipher import AES
import base64
import chardet
import json
import sys
if sys.stdout.encoding != 'UTF-8':
    sys.stdout.reconfigure(encoding='utf-8')  # Python 3.7+
# 在upload路由中添加
import logging
from datetime import datetime
import cv2
import numpy as np
# 配置数据库
# 新增依赖导入
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import text
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user
from werkzeug.security import generate_password_hash, check_password_hash
from sqlalchemy.orm import Session

# 配置 logging 基本格式（初始化时统一配置）
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    encoding="utf-8",  # 指定编码为 UTF-8
    handlers=[logging.FileHandler("app.log"), logging.StreamHandler()]
)


app = Flask(__name__)
# 获取当前脚本的绝对路径
base_dir = os.path.dirname(os.path.abspath(__file__))
# 构建相对路径
upload_folder = os.path.join(base_dir, 'uploads')
# 设置上传文件夹
app.config['UPLOAD_FOLDER'] = upload_folder
# 确保上传文件夹存在
if not os.path.exists(upload_folder):
    os.makedirs(upload_folder)
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 50MB限制
ALLOWED_EXTENSIONS = {'csv', 'txt', 'xlsx','jpg','peg','jpeg'}



# 初始化数据库和登录管理器
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://root:123456@localhost/datamasking?charset=utf8mb4'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.secret_key = os.urandom(24)  # 生产环境需使用固定密钥

db = SQLAlchemy(app)
login_manager = LoginManager(app)

# 用户模型
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(200), nullable=False)

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)

@login_manager.user_loader
def load_user(user_id):
    with Session(db.engine) as session:
        return session.get(User, int(user_id))


# 仅在开发环境初始化数据库
if os.environ.get('FLASK_ENV') == 'development':
    with app.app_context():
        db.create_all()


# 自动创建上传目录
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 敏感数据模式配置
# 正则表达式
SENSITIVE_PATTERNS = {
    "id_card": r'\b\d{15}\b|\b\d{17}[\dXx]\b', #身份证
    "phone": r'1[3-9]\d{9}', #电话
    "email": r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', #邮箱
    "bank_card": r'\b\d{16,19}\b', #银行卡账号
    "credit_card": r'\b\d{13,16}\b', #信用卡账号
    "salary": r'^¥?\d+￥$',
    "passport": r'\b[A-Za-z]\d{7,8}\b|\b[A-Z]{1,2}\d{6,9}\b',  # 中英文护照
    "ssn": r'\b\d{3}-\d{2}-\d{4}\b|\b\d{18}\b',              # 社保号
    "driver_license": r'\b[\u4e00-\u9fa5][A-Z0-9]{10,12}\b', # 驾驶证
    "name": r'^[\u4e00-\u9fa5]{2,4}$',  # 姓名通常为2-4个汉字
    "hire_date": r'\b\d{4}-\d{2}-\d{2}\b',  #雇佣日期
    "address": r'^[\u4e00-\u9fa5]+?(省|自治区|特别行政区) [\u4e00-\u9fa5]+?(市|地区) [\u4e00-\u9fa5]+?(县|区|镇|乡)$',
    "payment_password": r'\b\d{6}\b',                        # 支付密码
    "transaction_amount": r'\b¥?\d{1,3}(,\d{3})*(\.\d{2})?\b', # 交易金额
    "credit_limit": r'\b\d{1,3}(,\d{3})*\b',                 # 信用卡额度
    "company": r'\b.*(公司|集团|有限|股份|Company|Inc\.|Ltd\.).*\b', # 公司名称
    "annual_income": r'\b¥?\d{1,3}(,\d{3})*(\.\d{2})?\b',    # 年收入
    "employee_id": r'\bEMP\d{4,8}\b',                        # 工号
    "student_id": r'\bSTU\d{8,12}\b'                         # 学号
}
import string
class DataDesensitizer:
    """数据脱敏处理器"""
    def __init__(self, method='mask', secret_key=None):
        self.method = method
        self.secret_key = secret_key or os.urandom(16).hex()  # 默认随机密钥
        
        # 初始化加密组件
        if method == 'encrypt':
            self.cipher = AES.new(self.secret_key.encode(), AES.MODE_EAX)

    def process(self, value, data_type):
        """执行脱敏处理"""
        # 转换科学记数法为完整字符串
        if isinstance(value, float) and 'E+' in str(value):
              value = f"{int(value):d}"
       
        if pd.isna(value) or str(value).strip() == '':
            return value
            
        processor = getattr(self, f'_{self.method}', self._mask)
        return processor(str(value), data_type)

    # 不同脱敏策略实现
    def _mask(self, value, data_type):
        """默认遮盖策略"""
        rules = {
                'id_card': lambda x: x[:6] + "********" + x[-4:],
                'phone': lambda x: x[:3] + "****" + x[-4:],
                
                'email': lambda x: x.split("@")[0][:2] + "****@" + x.split("@")[1],
                'bank_card': lambda x: x[:6] + "*"*(len(x)-10) + x[-4:],
                'credit_card': lambda x: x[:4] + "*"*(len(x)-8) + x[-4:],
                "salary": lambda x: re.sub(r'\d', '*', x.split('￥')[0]) + '￥',
                'passport': lambda x: x[0] + "*"*(max(len(x)-2,0)) + x[-1:] if len(x)>1 else x,
                'ssn': lambda x: re.sub(r'\d', '*', x) if len(x)>5 else x,
                'driver_license': lambda x: x[:2] + "****" + x[-2:] if len(x)>4 else x,
                'name': lambda x: (x[0] + "*"*(len(x)-1)) if (len(x)>1 and re.match(r'^[\u4e00-\u9fa5]+$', x)) 
                                else (x.split()[0][0] + "*** " + x.split()[-1][0] if ' ' in x else x[0] + "***"),
                'address': lambda x: re.sub(r'(\d+)(号|栋|幢)', lambda m: '*'*len(m.group(1)) + m.group(2), x),
                'annual_income': lambda x: re.sub(r'\d', '*', x),
                'employee_id': lambda x: re.sub(r'\d', '*', x),
                'student_id': lambda x: x[:3] + "****" + x[-4:] if len(x)>7 else x
                
            }
        
        try:
            return rules.get(data_type, lambda x: x[:2] + "****" + x[-2:] if len(x)>4 else x)(value)
        except Exception as e:
            print(f"脱敏失败：{data_type} - {value} | 错误：{str(e)}")
            return value  # 返回原始值或根据业务需求处理


    def _hash(self, value, data_type):
        """哈希脱敏（加盐）"""
        salt = os.urandom(16).hex()
        return hashlib.pbkdf2_hmac('sha256', value.encode(), salt.encode(), 100000).hex()

    def _replace(self, value, data_type):
        """修正的替换逻辑"""
        print(f"Replacing {data_type}: {value}")  # 记录日志，方便排查问题
        if data_type == 'phone':
            return f'1{random.randint(3,9)}{random.randint(0,9)}{random.randint(10000000,99999999)}'
        elif data_type == 'salary':
            return '3000￥'
        elif data_type == 'email':
            prefix = ''.join(random.choices(string.ascii_letters + string.digits, k=8)).lower()
            return f'{prefix}@tuoming.com'
        elif data_type == 'id_card':
         def _checksum(id17):
                weight = [7,9,10,5,8,4,2,1,6,3,7,9,10,5,8,4,2]
                check_code = '10X98765432'
                total = sum(int(id17[i]) * weight[i] for i in range(17))
                return check_code[total % 11]

         address_code = value[:6] if len(value)>=6 else '110101'
         birth_date = f"{random.randint(1960,2023):04d}{random.randint(1,12):02d}{random.randint(1,28):02d}"
         sequence = f"{random.randint(0,999):03d}"
         checksum = _checksum(address_code + birth_date + sequence)
         return f"{address_code}{birth_date}{sequence}{checksum}"
        
        elif data_type == 'hire_date':
            # 修正 hire_date 生成逻辑，避免变成地址
            return f"{random.randint(2000,2023)}-{random.randint(1,12):02d}-{random.randint(1,28):02d}"

        elif data_type == 'address':
            provinces = ['吉林省', '陕西省', '广西壮族自治区', '江西省', '西藏自治区', '台湾省', '香港特别行政区']
            cities = {
                '吉林省': ['吉林市', '通化市', '白城市', '延边州'],
                '陕西省': ['安康市', '汉中市', '榆林市', '宝鸡市'],
                '广西壮族自治区': ['防城港市', '柳州市', '桂林市', '北海市'],
                '江西省': ['青安市', '赣州市', '宜春市', '九江市'],
                '西藏自治区': ['阿里地区', '日喀则市', '林芝市', '山南市'],
                '台湾省': ['菊果县', '新竹市', '彰化县'],
                '香港特别行政区': ['新界', '九龙', '香港岛']
            }
            counties = ['永吉县', '镇坪县', '上思县', '吉安县', '喝尔县', '泄安乡', '集安县', '观塘区', '宁强县', '融安县', '南康区']

            province = random.choice(provinces)
            city = random.choice(cities[province])
            county = random.choice(counties)
            return f"{province} {city} {county}"

        elif data_type == 'passport':
            return random.choice(['E','G']) + str(random.randint(10000000, 99999999))
        elif data_type == 'ssn':
         return f'{random.randint(100,999)}-{random.randint(10,99)}-{random.randint(1000,9999)}'
        elif data_type == 'name':
            surnames = ['王','李','张','刘','陈','赵','黄','周']
            given_names = ['伟','芳','娜','强','敏','杰','婷','浩']
            return f'{random.choice(surnames)}{random.choice(given_names)}'
        
        elif data_type == 'company':
            suffixes = ['科技有限公司', '集团有限公司', '股份有限公司']
            return f'{random.choice(["东方","华腾","新创"])}{random.choice(suffixes)}'
        
        elif data_type == 'employee_id':
            return f'EMP{random.randint(20230000, 20239999)}'
        else:
            return value  # 其他类型保持不变

    def _encrypt(self, value, data_type):
        """AES加密脱敏"""
        nonce = os.urandom(16)
        cipher = AES.new(self.secret_key.encode(), AES.MODE_EAX, nonce=nonce)
        ciphertext, tag = cipher.encrypt_and_digest(value.encode())
        return base64.b64encode(nonce + tag + ciphertext).decode()

def desensitize_image(filepath):
    """图像脱敏处理 - 人脸检测和模糊"""
    try:
        # 读取图像
        image = cv2.imread(filepath)
        if image is None:
            raise ValueError("无法读取人脸图像文件")
        
        # 加载人脸检测器
        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        
        # 转换为灰度图像
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 检测人脸
        faces = face_cascade.detectMultiScale(
            gray, 
            scaleFactor=1.1, 
            minNeighbors=5, 
            minSize=(30, 30),
            flags=cv2.CASCADE_SCALE_IMAGE
        )
        
        # 对检测到的人脸进行模糊处理
        for (x, y, w, h) in faces:
            face_roi = image[y:y+h, x:x+w]
            blurred_face = cv2.GaussianBlur(face_roi, (99, 99), 30)
            image[y:y+h, x:x+w] = blurred_face
        
        # 保存处理后的图像
        output_filename = f"processed_{os.path.splitext(os.path.basename(filepath))[0]}.jpg"
        output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
        cv2.imwrite(output_path, image)
        
        return output_filename
    except Exception as e:
        logging.error(f"人脸图像脱敏失败: {str(e)}")
        raise

def detect_sensitive_columns(df):
    """检测包含敏感数据的列"""
    sensitive_columns = {}
    for col in df.columns:
        for data_type, pattern in SENSITIVE_PATTERNS.items():
            try:
                if df[col].astype(str).str.contains(pattern, regex=True, na=False).any():
                    sensitive_columns[col] = data_type
                    break
            except Exception as e:
                logging.error(f"Error processing column {col} with pattern {pattern}: {e}")
    return sensitive_columns

@app.before_request
def check_auth():
    allowed_endpoints = ['auth', 'static']
    if request.endpoint in allowed_endpoints:
        return
    if not current_user.is_authenticated:
        return redirect(url_for('auth'))  # 



@app.route('/auth', methods=['GET', 'POST'])
def auth():
    """统一处理登录/注册"""
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        action = request.form.get('action')  # login 或 register

        # 注册逻辑
        if action == 'register':
            existing_user = User.query.filter_by(username=username).first()
            if existing_user:
                flash('用户名已存在')
                return redirect(url_for('auth'))
            
            new_user = User(username=username)
            new_user.set_password(password)
            db.session.add(new_user)
            db.session.commit()
            flash('注册成功，请登录')
            return redirect(url_for('auth'))

        # 登录逻辑
        elif action == 'login':
            user = User.query.filter_by(username=username).first()
            if user and user.check_password(password):
                login_user(user)
                return redirect(url_for('index'))
            else:
                flash('用户名或密码错误')

    return render_template('auth.html')



@app.route('/index')
@login_required
def index():
    """登录后的系统主页"""
    return render_template('index.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('auth'))

@app.route('/upload', methods=['POST'])
@login_required
def upload():
    if 'file' not in request.files:
        return "请选择要上传的文件"
    
    file = request.files['file']
    method = request.form.get('method', 'mask')
    
    if not file or file.filename == '':
        return "未选择文件"
    
    if not allowed_file(file.filename):
        return "不支持的文件格式"

    filename = secure_filename(file.filename)
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(filepath)

    try:
        # 处理图片文件
        if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
            if method == 'image_mask':
                output_filename = desensitize_image(filepath)
                report = {
                    'method': 'IMAGE_MASK',
                    'original_file': filename,
                    'processed_file': output_filename,
                    'sensitive_columns': 1,
                    'secret_key': None
                }
                
                 # 在文件处理成功后添加
                log_process(
                    user_id=current_user.id,
                    filename=secure_filename(filename),
                    method=request.form.get('method'),
                    encrypt_key=None
                )
                
                return render_template('result.html', report=report)
            else:
                return "请选择人脸图像遮盖脱敏方法处理图片文件"

        # 检测文件编码
        with open(filepath, 'rb') as f:
            encoding = chardet.detect(f.read())['encoding'] or 'utf-8'
        
        # 读取文件
        if filename.lower().endswith('.xlsx'):
            df = pd.read_excel(filepath, dtype=str)
        else:
            df = pd.read_csv(filepath, 
                           engine='python', 
                           dtype=str, 
                           encoding=encoding,
                           on_bad_lines='warn')

        # 执行脱敏处理
        desensitizer = DataDesensitizer(method=method)
        sensitive_cols = detect_sensitive_columns(df)
        
        for col, data_type in sensitive_cols.items():
            df[col] = df[col].apply(
                lambda x: desensitizer.process(x, data_type)
            )

        # 保存处理结果
        output_filename = f"processed_{os.path.splitext(filename)[0]}.csv"
        output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
        df.to_csv(output_path, index=False, encoding='utf-8-sig')

        # filename = filename.encode("unicode_escape").decode("utf-8")  # 转义特殊字符
        logging.info(f"用户 {request.remote_addr} 于 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} 处理了文件 {filename}")

        # 生成报告
        report = {
            'method': method.upper(),
            'original_file': filename,
            'processed_file': output_filename,
            'sensitive_columns': len(sensitive_cols),
            'secret_key': desensitizer.secret_key if method == 'encrypt' else None
        }

        # 记录处理日志
        log_process(
            user_id=current_user.id,
            filename=filename,
            method=method,
            encrypt_key=desensitizer.secret_key if method == 'encrypt' else None
        )

        return render_template('result.html', report=report)
    
    except Exception as e:
        return f"处理失败: {str(e)}"
    finally:
        if os.path.exists(filepath):
            os.remove(filepath)
    
# 记录脱敏操作（在upload路由成功后调用）
def log_process(user_id, filename, method, encrypt_key=None):
    """记录处理日志"""
    try:
        with app.app_context():
            # 使用text()包裹SQL语句
            db.session.execute(
                text("""
                    INSERT INTO process_records 
                    (user_id, filename, method, encrypt_key) 
                    VALUES (:uid, :fn, :mt, :key)
                """),
                {
                    'uid': user_id,
                    'fn': filename,
                    'mt': method,
                    'key': encrypt_key
                }
            )
            db.session.commit()
    except Exception as e:
        logging.error(f"数据库写入失败: {str(e)}")
        raise

@app.route('/records')
@login_required
def view_records():
    """查看处理记录"""
    try:
        # 使用text()包裹SQL语句
        records = db.session.execute(
            text("""
                SELECT filename, method, encrypt_key, process_time 
                FROM process_records 
                WHERE user_id = :uid 
                ORDER BY process_time DESC
            """),
            {'uid': current_user.id}
        ).fetchall()
        return render_template('record.html', records=records)
    except Exception as e:
        logging.error(f"查询处理记录失败: {str(e)}")
        return "无法获取处理记录", 500
    
@app.route('/decrypt', methods=['GET','POST'])
@login_required
def decrypt_data():
    secret_key = request.form.get("key", "")
    print(f"[DEBUG] 用户输入的密钥：{secret_key}")  # 在控制台查看实际密钥
    if request.method == 'GET':
        return render_template('decrypt.html')  # 显示解密页面
    if 'file' not in request.files:
        return "请上传加密文件"
    
    file = request.files['file']
    secret_key = request.form.get('key', '')
    
    # 保存临时文件
    temp_path = os.path.join(app.config['UPLOAD_FOLDER'], 'temp_encrypted.csv')
    file.save(temp_path)
    
    # 执行解密
    output_path = os.path.join(app.config['UPLOAD_FOLDER'], 'decrypted.csv')
    decrypt_file(temp_path, output_path, secret_key)
    
    return send_file(output_path, as_attachment=True)


@app.route('/desensitize', methods=['GET', 'POST'])
def desensitize_page():
    if request.method == 'GET':
        return render_template('desensitize.html')  # 新的脱敏配置页面
    
    # POST方法处理带参数的脱敏请求
    method = request.form.get('method')
    salt = request.form.get('salt', '')
    # ... 其他参数处理 ...

@app.route('/api/v1/process', methods=['POST'])
@login_required
def api_process():
    """
    RESTful API 数据处理接口
    请求格式：
    {
        "data": [
            {"phone": "13800138000", "id_card": "110105199901011234", ...},
            ...
        ],
        "method": "mask|hash|replace|encrypt",
        "encrypt": true|false
    }
    响应格式：
    {
        "status": "success|error",
        "data": [处理后的数据],
        "secret_key": "加密密钥（如果启用加密）",
        "error": "错误信息（如果失败）"
    }
    """
    try:
        # 验证请求格式
        if not request.is_json:
            return jsonify({"status": "error", "error": "仅支持JSON格式请求"}), 400

        req_data = request.get_json()
        
        # 参数验证
        required_fields = ['data', 'method']
        if any(field not in req_data for field in required_fields):
            return jsonify({"status": "error", "error": "缺少必要参数"}), 400

        # 创建DataFrame
        # df = pd.DataFrame(req_data['data'])
        # 强制所有输入数据转为字符串
        df = pd.DataFrame(req_data['data']).astype(str)
        
        # 执行脱敏处理
        desensitizer = DataDesensitizer(
            method=req_data.get('method', 'mask'),
            secret_key=os.urandom(16).hex() if req_data.get('encrypt') else None
        )
        
        # 检测敏感列
        sensitive_cols = detect_sensitive_columns(df)
        
        # 处理数据
        for col, data_type in sensitive_cols.items():
            df[col] = df[col].apply(
                lambda x: desensitizer.process(x, data_type)
            )

        # 构建响应数据
        result = {
            "status": "success",
            "data": df.to_dict(orient='records')
        }
        

        # 添加加密密钥
        if req_data.get('encrypt'):
            result["secret_key"] = desensitizer.secret_key

        return jsonify(result), 200

    except Exception as e:
        return jsonify({
            "status": "error",
            "error": f"处理失败: {str(e)}"
        }), 500


@app.route('/download/<filename>')
@login_required
def download(filename):
    path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if os.path.exists(path):
        return send_file(path, as_attachment=True)
    return "文件不存在"


@app.route('/download_key')
@login_required
def download_key():
    key = request.args.get('key')
    # 这里根据实际需求处理密钥下载逻辑，例如返回密钥文件
    # 简单示例：
    return key

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)